import re

# 计算缩进数量
def nsep(line, sep=' '):
    for i in range(len(line)):
        if line[i] != sep: return i
    return len(line)

# 解析字符串，注释，多行注释等。
def common(line, cur=None):
    keys = ['"""', "'''", '"', "'", '#']
    if not cur is None: keys = [cur]
    rst, start = None, 1e4
    for c in keys:
        if not c in line: continue
        idx = line.index(c)
        if idx < start: rst, start = c, idx
    if cur==None and rst==None:
        return line, '', 'ok'
    if cur==None and rst=='#':
        return line[:start], '', 'ok'
    if cur==None and rst in keys:
        return line[:start], line[start+len(rst):], rst
    if cur in keys and rst==cur:
        return 'str', line[start+len(rst):], None
    return None, '', cur

# 去注释，字符串转化
def rm_common(lines):
    buf, rst, cur = [], [], None
    for line in lines:
        while True:
            line1, line, cur = common(line, cur)
            if not line1 is None: buf.append(line1)
            if line1 is None or cur == 'ok': break
        if cur == 'ok':
            cur = None
            rst.append(''.join(buf))
            buf = []
        else: rst.append('')
    return rst

# 解析括号
def bracket(line, cur=''):
    keys = '{[()]}'
    rst, start = '', 1e4
    for c in keys:
        if not c in line: continue
        idx = line.index(c)
        if idx < start: rst, start = c, idx

    if rst == '': return line, '', cur
    if rst in '{[(':
        return line[:start+1], line[start+1:], cur+rst
    if rst in ')]}' and cur[-1:]+rst in {'()','[]','{}'}:
        return line[:start+1], line[start+1:], cur[:-1]
    if rst in ')]}' and cur[-1:]+rst not in {'()','[]','{}'}:
        return line[:start+1], line[start+1:], None
    return line, None, ''

# 分号拆分多行
def spilit_line(line):
    ls = line.split(';')
    sep = nsep(line)
    return [' '*sep + i.strip() for i in ls]

# 去括号
def rm_bracket(lines):
    buf, rst, cur, sep = [], [], None, 0
    for line in lines:
        lbuf = []
        for code in spilit_line(line):
            if nsep(code) <= sep:
                buf, cur = [], ''
            if cur=='': sep = nsep(code)
            while True:
                code1, code, cur = bracket(code, cur)
                # print(code1, '|||', code, '|||', cur)
                buf.append(code1)
                if code=='': break
            if cur=='': lbuf.append(''.join(buf))
            else: lbuf.append(' '*sep)
            if cur in {None, ''}: buf, cur = [], ''
        rst.append(';'.join(lbuf))
    return rst

# 将数字替换为num
def rm_num(lines):
    rst = []
    for line in lines:
        line = re.sub(r'\d+\.\d+', 'num', line)
        f = lambda x:x.group(0)[:1]+'num'
        rst.append(re.sub(r'^\d+|\W\d+', f, line))
    return rst

# 将容器内的元素去掉，将函数调用参数去掉
def rm_para(lines):
    buf, rst, cur, sep = [], [], None, 0
    for line in lines:
        lbuf = []
        for code in spilit_line(line):
            if 'def ' in line or 'class' in line:
                lbuf.append(code); continue
            buf, cur = [], ''
            for i in code:
                if cur[-1:]+i in {'{}','[]','()'}:
                    cur = cur[:-1];
                if cur=='': buf.append(i)
                if i in '[{': cur += i
                if len(buf)>1 and i=='(':
                    if re.match('[\w|\]|\)]', buf[-2]): cur += i
                if i==')' and ',' in buf:
                    if buf[::-1].index(',')<buf[::-1].index('('):
                        del buf[-buf[::-1].index('('):-1]
            lbuf.append(''.join(buf))
        rst.append(';'.join(lbuf))
    return rst

if __name__ == '__main__':
    cont = '''
    a = 1 # xxxxxx
    b = 'abc # xxxxxx'
    """
    'abc',"def" # hahaha
    """
    a, b = """123""", '456'
    
    a = ((1+3)* 4 + 'abc')
    a = 2
    c = 1
    a = 1; b=2; c = [2,
        "abc"]
    d = {'a':1}
    t = ([1,2,3], 2)
    f8(12)
    def f(a,
        d = 1):
        return 123

    rst = f(1,2,3)

    import numpy as np
    arr = np.array([1])
    '''
    
    lines = cont.split('\n')
    rst = rm_common(lines)

    for a, b in zip(lines, rst):
        print('%-30s | %-30s'%(a, b))

    print('-'*60)
    lines = rst
    rst = rm_bracket(lines)

    for a, b in zip(lines, rst):
        print('%-30s | %-30s'%(a, b))

    print('-'*60)
    lines = rst
    rst = rm_num(lines)

    for a, b in zip(lines, rst):
        print('%-30s | %-30s'%(a, b))

    print('-'*60)
    lines = rst
    rst = rm_para(lines)

    for a, b in zip(lines, rst):
        print('%-30s | %-30s'%(a, b))
